{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "69f38f18-cf7b-4ba8-b707-52c200c34fcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision as tv\n",
    "import torch.nn as nn\n",
    "import torchvision.transforms as transforms\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "import torch.utils.data as Data\n",
    "import sklearn.metrics as metrics\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "import os\n",
    "import os.path as osp\n",
    "import matplotlib.pyplot as plt\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6e384268-88c3-42ca-ad64-761cf8939f39",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = \"./data/road/\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45607ec8-f095-4799-8ccc-6a30a0e21987",
   "metadata": {},
   "source": [
    "# Load Efficientnet from network "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "53928672-524b-48a2-8ae5-e25c86f37383",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/avicii/miniconda3/envs/DL/lib/python3.11/site-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.\n",
      "  warnings.warn(\n",
      "/home/avicii/miniconda3/envs/DL/lib/python3.11/site-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=EfficientNet_B0_Weights.IMAGENET1K_V1`. You can also use `weights=EfficientNet_B0_Weights.DEFAULT` to get the most up-to-date weights.\n",
      "  warnings.warn(msg)\n",
      "Downloading: \"https://download.pytorch.org/models/efficientnet_b0_rwightman-7f5810bc.pth\" to /home/avicii/.cache/torch/hub/checkpoints/efficientnet_b0_rwightman-7f5810bc.pth\n",
      "100%|██████████| 20.5M/20.5M [00:04<00:00, 4.82MB/s]\n"
     ]
    }
   ],
   "source": [
    "model = tv.models.efficientnet_b0(pretrained=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "9c7e2c05-38c6-446b-af2e-967959aceb8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.classifier[1] = nn.Sequential(\n",
    "    nn.Linear(model.classifier[1].in_features, 1, bias=True),\n",
    "    nn.Sigmoid()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36450bc8-65d6-4860-b257-4527174e38b3",
   "metadata": {},
   "source": [
    "# define a function to read pictures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "d2a6bce2-455e-45b0-ab13-6cdced9f7903",
   "metadata": {},
   "outputs": [],
   "source": [
    "trans = transforms.Compose([\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Resize([224, 224])\n",
    "        ])\n",
    "def read_pics(img_raw_path):\n",
    "    target = []\n",
    "    imgs = torch.tensor([])\n",
    "    for f in os.listdir(img_raw_path):\n",
    "        f = osp.join(img_raw_path, f)\n",
    "        pn, fn = osp.split(f)\n",
    "        target.append(1 if \"normal\" in fn else 0)\n",
    "        try:\n",
    "            img = plt.imread(f)\n",
    "        except:\n",
    "            continue\n",
    "        img = trans(img).unsqueeze(0)\n",
    "        imgs = torch.cat((img, imgs if len(imgs) != 0 else torch.tensor([])), dim = 0)\n",
    "    return imgs, torch.tensor(target).type(torch.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "59e98cb9-aa06-4b34-930b-2f6c65770dc8",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/avicii/miniconda3/envs/DL/lib/python3.11/site-packages/torchvision/transforms/functional.py:1603: UserWarning: The default value of the antialias parameter of all the resizing transforms (Resize(), RandomResizedCrop(), etc.) will change from None to True in v0.17, in order to be consistent across the PIL and Tensor backends. To suppress this warning, directly pass antialias=True (recommended, future default), antialias=None (current default, which means False for Tensors and True for PIL), or antialias=False (only works on Tensors - PIL will still use antialiasing). This also applies if you are using the inference transforms from the models weights: update the call to weights.transforms(antialias=True).\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "X = read_pics(data_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "b35c8d94-6fa6-4134-8774-c1802bc82ac2",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch_dataset  = Data.TensorDataset(X, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "700019d1-3bce-469a-bdf2-d11eb1ec15e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "loader = Data.DataLoader(dataset=torch_dataset, shuffle=True, batch_size = 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "ced8776c-7ed5-4df4-8fab-5f20a9db2971",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss = nn.BCELoss()\n",
    "lr = 1e-5\n",
    "params_1x = [param for name, param in model.named_parameters()\n",
    "                 if name not in [\"classifier.1.0.weight\", \"classifier.1.0.bias\"]]\n",
    "trainer = torch.optim.SGD([{'params': params_1x},\n",
    "                                  {'params':model.classifier[1].parameters(),\n",
    "                                  'lr': lr * 10}],lr = lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "c7bffd62-2b76-4744-b6b6-36b47a7d9b0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(X, y, loader, net,epoch,trainer, batch_size):  \n",
    "        loss_func = nn.BCELoss()\n",
    "        net.train()\n",
    "        for i in range(epoch):\n",
    "            for step, (batch_x, batch_y) in enumerate(loader):\n",
    "                trainer.zero_grad()\n",
    "                output = net(batch_x)\n",
    "                loss = loss_func(output[:, 0], batch_y.type(torch.float32))\n",
    "                loss.backward()\n",
    "                trainer.step()\n",
    "            \n",
    "        net.eval()\n",
    "        for i in np.arange(0.1, 0.6, 0.1):\n",
    "            pred = []\n",
    "            for x in X:\n",
    "                y_hat = net(x.unsqueeze(0))\n",
    "        #             _, ax = plt.subplots(1,1)\n",
    "        #             print(y_hat.detach().numpy()[0])\n",
    "        #             ax.imshow(np.transpose(X[idx], (1,2,0)))\n",
    "        #             plt.show()\n",
    "                y_pred = 1 if y_hat.detach().numpy()[0] > i else 0\n",
    "                pred.append(y_pred)\n",
    "            acc = metrics.accuracy_score(y.detach().numpy(), pred)\n",
    "            f1 = metrics.f1_score(y.detach().numpy(), pred)\n",
    "            precision = metrics.precision_score(y.detach().numpy(), pred)\n",
    "            recall = metrics.recall_score(y.detach().numpy(), pred)\n",
    "            print(f'threshold = {i} , acc = {acc}, f1 = {f1}, precision = {precision}, recall = {recall}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "3f40e25e-81ab-43dc-8821-ff4794ccc74f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "threshold = 0.1 , acc = 0.8837209302325582, f1 = 0.9382716049382717, precision = 0.8837209302325582, recall = 1.0\n",
      "threshold = 0.2 , acc = 0.8837209302325582, f1 = 0.9382716049382717, precision = 0.8837209302325582, recall = 1.0\n",
      "threshold = 0.30000000000000004 , acc = 0.8837209302325582, f1 = 0.9382716049382717, precision = 0.8837209302325582, recall = 1.0\n",
      "threshold = 0.4 , acc = 0.8837209302325582, f1 = 0.9382716049382717, precision = 0.8837209302325582, recall = 1.0\n",
      "threshold = 0.5 , acc = 0.8571428571428571, f1 = 0.923076923076923, precision = 0.8805460750853242, recall = 0.9699248120300752\n"
     ]
    }
   ],
   "source": [
    "train(X, y, loader, model, 1, trainer, 4)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
